
import sdl2
from sketchz import *
from random import randint, choice
from itertools import product

# ~ window = WindowAndroid((600, 800))
window = Window((600, 800))

window.img.clear_value = (0.6, 0.7, 0.8, 1)





prog1 = {
    "vertex_shader":'''
#version 320 es
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec3 aNormal;

layout (std140) uniform Common {
    mat4 modelMat;
    mat4 viewMat;
    mat4 projMat;
    vec3 lightPos;
    float n1;
    vec3 viewPos;
    float n2;
    vec3 objectColor;
    float n3;
};

out vec3 mPos;
out vec3 wPos;
out vec3 vPos;
out vec3 wNormal;
out vec3 vNormal;

void main() {
    mPos = aPos;
    vec4 wPos4 = modelMat * vec4(aPos, 1.0);
    vec4 vPos4 = viewMat * wPos4;
    wPos = wPos4.xyz;
    vPos = vPos4.xyz;
    wNormal = mat3(modelMat) * aNormal;
    vNormal = mat3(viewMat * modelMat) * aNormal;
    gl_Position = projMat * vPos4;
  }
            ''',
    "fragment_shader":'''
#version 320 es
precision highp float;

in vec3 mPos;
in vec3 wPos;
in vec3 vPos;
in vec3 wNormal;
in vec3 vNormal;

layout (std140) uniform Common {
    mat4 modelMat;
    mat4 viewMat;
    mat4 projMat;
    vec3 lightPos;
    float n1;
    vec3 viewPos;
    float n2;
    vec3 objectColor;
    float n3;
};

layout (location = 0) out vec4 FragColor;
layout (location = 1) out vec4 Help;

float orf(float a, float b)
{
    return 1.0 - (1.0 - a) * (1.0 - b);
}

float cross2(vec2 a, vec2 b)
{
    return a.x * b.y - a.y * b.x;
}
float dot2(vec2 a, vec2 b)
{
    return a.x * b.x + a.y * b.y;
}

float hill(float a, float b, float x)
{
    return step(a, x) - step(b, x);
}

float piece(float x0, float y0, vec2 xy)
{
    float pi = 3.14159;
    vec2 b1 = vec2(0.3 + x0, y0);
    vec2 b2 = vec2(-0.3 + x0, y0);
    float tc = dot2(normalize(b2-xy), normalize(b1-xy));
    tc = acos(tc);

    float board = hill(-0.3 +x0, 0.3+x0, xy.x) * hill(y0 +0.0, y0+1.0, xy.y);
    float k = 0.5;
    if (k > 0.0)
    {
       board = board * (1.0 - smoothstep(pi - k- 0.05, pi-k+0.05, tc));
    }
    else
    {
        board = orf(board, smoothstep(pi + k- 0.05, pi+k+ 0.05, tc));
    }
    return board;
}


void main()
{
    // gamma jiao4zheng4
    vec3 objectColor_line = objectColor;

    float pi = 3.14159;
    float hpi = pi / 2.0;

    vec3 fd = vec3(0.0, 0.0, 1.0);
    vec3 up = vec3(0.0, 1.0, 0.0);
    vec3 lt = vec3(1.0, 0.0, 0.0);
    //float y = acos(dot(mPos, up));
    vec3 temp = normalize(cross(up, mPos)); // self left
    vec3 temp2 = cross(lt, temp) ;
    float x = asin(length(temp2)*sign(dot(temp2, up)));
    vec3 temp3 = normalize(cross(temp, up)) ; // self fd
    vec3 temp4 = cross(temp3, mPos) ;
    float y = asin(length(temp4)*sign(dot(temp4, -lt)));
    vec2 xy = vec2(x, y);

    float basey = sin(n3 * 5.0) * 0.1 + 0.1;
    float basey2 = sin(n3 * 5.0 +1.0) * 0.1 + 0.1;

    float bo1 = piece(-0.45, basey, xy);
    float bo2 = piece(0.45, basey2, xy);
    float board =orf(bo1, bo2);

    float r = length(vec2((x-0.45) * 2.2, y*1.3));
    r = 1.0-smoothstep(0.5, 0.55, r);
    float r2 = length(vec2((x+0.45) * 2.2, y*1.3));
    r2 = 1.0-smoothstep(0.5, 0.55, r2);
    objectColor_line.b = orf(r, r2);
    objectColor_line.b = objectColor_line.b*(1.0-board);
      //  objectColor_line.r = board;
    objectColor_line.r = objectColor_line.b;
    objectColor_line.g = objectColor_line.b;


    objectColor_line.r = pow(objectColor_line.r, 2.0);
    objectColor_line.g = pow(objectColor_line.g, 2.0);
    objectColor_line.b = pow(objectColor_line.b, 2.0);

    float ambientStrength = 0.05;
    vec3 ambient = ambientStrength * objectColor_line;

    vec3 norm = normalize(wNormal);
    vec3 lightDir = normalize(lightPos - wPos);
    float diff = max(dot(norm, lightDir) , 0.0);
    //vec3 diffuse = diff * vec3(0.5, 0.5, 0.5) ;
    vec3 diffuse = diff * objectColor_line ;

    vec3 viewDir = normalize(viewPos - wPos);
    float spec = 0.0;
    vec3 halfwayDir = normalize(lightDir + viewDir);
    spec = pow(max(dot(wNormal, halfwayDir), 0.0), 32.0);
    vec3 specular = spec * vec3(0.3);

    vec3 result = (ambient + diffuse + specular * 0.1);
    result.r = pow(result.r, 0.5);
    result.g = pow(result.g, 0.5);
    result.b = pow(result.b, 0.5);

    //vec3 red = vec3(1, 0, 0);

    FragColor = vec4(result, 1.0);

    //Help = vec4(normalize(vNormal) , length(vPos));
    Help = vec4(normalize(vNormal)  * length(vPos), 1);


}
            '''
}


class Obj(Drawable, Moveable):
    def __init__(self):
        self.vertices = []
        self.normals = []
        self.color = V3(0.7, 0.7, 0.7)
        self.prog = prog1

        self.t = Turtle()
        self.box = None

        self.pipeline = None
        self.uniform_buffer = None

    def u(self, window):       # update
        assert len(self.vertices) != 0
        assert hasattr(window, "render")
        assert hasattr(window, "update")
        assert hasattr(window, "mainloop")

        vl = np.array(self.vertices, dtype="f4")
        nl = np.array(self.normals, dtype="f4")
        vn = np.dstack([vl[:, 0], vl[:, 1], vl[:, 2], nl[:, 0], nl[:, 1], nl[:, 2]])
        buf = window.ctx.buffer(vn.tobytes())
        try:
            self.uniform_buffer = window.ctx.buffer(size=64+64+64+16+16+16, dynamic=True)
        except:
            self.uniform_buffer = window.ctx.buffer(size=64+64+64+16+16+16, uniform=True)
        self.pipeline = window.ctx.pipeline(
            vertex_shader = self.prog["vertex_shader"],
            fragment_shader = self.prog["fragment_shader"],
            framebuffer = [window.img, window.imgh, window.imgs],
            layout=[
                {'name': 'Common', 'binding': 0},
            ],
            resources=[
                {
                    'type': 'uniform_buffer',
                    'binding': 0,
                    'buffer': self.uniform_buffer,
                },
            ],
            vertex_buffers = zengl.bind(buf, "3f 3f", 0, 1),
            vertex_count = buf.size // zengl.calcsize('3f 3f'),
            # ~ blend = {'enable': True},
        )



    def copy(self):
        o = Obj()
        o.vertices = self.vertices.copy()
        o.normals = self.normals.copy()
        o.color = self.color
        o.prog = self.prog

        o.t = self.t.copy()
        if self.box is not None:
            o.box = self.box.copy(o)

        o.pipeline = self.pipeline
        o.uniform_buffer = self.uniform_buffer
        return o

    def as_unit(self):
        v = np.array(self.vertices) @ np.array(self.t.abs().rot.matrix33).T + self.t.abs().pos
        n = np.array(self.normals) @ np.array(self.t.abs().rot.matrix33).T
        return v, n

    def __rrshift__(self, obj):
        if len(obj) != 2:
            raise Exception(f"add part faild, {obj}")
        else:
            v, n = obj
            assert len(v) == len(n)
            self.vertices.extend(v)
            self.normals.extend(n)

    def m(self):
        t = self.t.abs()
        return ( Matrix44.from_translation(t.pos) * t.rot.matrix44.T).astype("f4")

    def draw(self, camera):
        if self.pipeline is None:
            print("没准备好绘制", self)
            return

        m = self.m()
        v = camera.v()
        p = camera.p()
        pos = camera.t.pos.astype("f4")

        self.uniform_buffer.write(struct.pack('64s64s64s3f1f3f1f3f1f', m.tobytes(), v.tobytes(), p.tobytes(), *pos, 0, *pos, 0, *self.color, G.t))
        self.pipeline.render()

    def pile(self, t, anchor:"U(-1~1, -1~1, -1~1)"):
        a = anchor * self.box.size / 2
        t = t.abs()
        self.t.pos = t.pos - t.rot * a
        self.t.rot= t.rot


def mid(l):
    return sum(l, U3(0, 0, 0)) / len(l)

def lerpf(b, k):
    return lambda x: lerp(x, b, k)

def take(l, *il):
    return [l[i] for i in il]
# 5 4 3
# 0 1 2



up = U3(0, 0, 1)
dn = U3(0, 0, -1)
fd = U3(0, 1, 0)
bk = U3(0, -1, 0)
lt = U3(-1, 0, 0)
rt = U3(1, 0, 0)

acc = [
up, lt, bk,
up, bk, rt,
dn, bk, lt,
dn, rt, bk,
up, fd, lt,
up, rt, fd,
dn, lt, fd,
dn, fd, rt,
]

def better(l):
    te = []
    for i in range(0, len(l), 3):
        te.append(l[i:i+3])
    new = []
    for p1, p2, p3 in te:
        p12 = lerp(p1, p2, 0.5).normalized
        p23 = lerp(p2, p3, 0.5).normalized
        p31 = lerp(p1, p3, 0.5).normalized
        new.extend([
        p1, p12, p31,
        p2, p23, p12,
        p3, p31, p23,
        p12, p23, p31,
        ])
    return new


acc = better(acc)
acc = better(acc)
acc = better(acc)
#




box = Obj()
box.color = (1, 1, 1)

(acc, easy_norm(acc)) >> box
box.u(window)
box >> window

camera = window.camera
camera.t.pos = U3(7, 7, 7)

walker = Viewer(camera)
walker >> window

class G:
    t = 0

class Ctrl(Control):
    def update(self, window, dt):
        G.t += dt
    def reg_event(self, window):
        window.update_callback.append(self.update)

Ctrl() >> window

window.mainloop()





